import torch
import torch.nn as nn
import network.metrics as met
from utils.label_transform import label2real
import numpy as np
from utils.cal_voxel_mse import init_vertex


def cal_voxel_mse_loss_tensor(vertex_size, interval_num, rt1, rt2, mode='average', weight=1):
    # 根据体素间距，生成vm的间隔采样稀疏点云
    cube = init_vertex(vertex_size, interval_num)
    cube = torch.from_numpy(cube).to(torch.float)
    new_pos1 = (rt1 @ cube.T).T
    new_pos2 = (rt2 @ cube.T).T
    # 计算两个新位置之间顶点的距离
    distance = torch.linalg.norm((new_pos1 - new_pos2), axis=1)
    distance = distance * weight
    # 给旋转中心增加权重
    if mode == 'average':
        return torch.mean(distance)
    else:
        return torch.max(distance)


def add_noise_tensor(angle_noise=None, trans_noise=None, rot_cen=None, show_info=False):
    """
    全tensor版本的添加噪声方式
    :param angle_noise: 角度噪声
    :param trans_noise: 位移噪声
    :param rot_cen: 旋转中心
    :param show_info: 是否显示过程（默认不显示）
    :return:
    """
    # 从体素到旋转中心
    ct2rot_cen = np.diag([1, 1, 1, 1]).astype(np.float32)
    ct2rot_cen[:3, 3] = -rot_cen
    ct2rot_cen = torch.from_numpy(ct2rot_cen)
    # 从旋转中心到加噪后的旋转中心
    # 1 添加旋转噪声
    pi = torch.from_numpy(np.pi)
    R = _euler_angles2rot_matrix_tensor(angle_noise[0] * pi / 180,
                                        angle_noise[1] * pi / 180,
                                        angle_noise[2] * pi / 180)
    # 添加平移噪声
    t_noise = torch.hstack((trans_noise, torch.tensor([1])))
    rot_cen_2noised_cen1 = torch.from_numpy(np.diag([1, 1, 1, 1]).astype(np.float32))
    rot_cen_2noised_cen2 = torch.from_numpy(np.diag([1, 1, 1, 1]).astype(np.float32))
    rot_cen_2noised_cen1[:3, :3] = R
    rot_cen_2noised_cen2[:, 3] = t_noise
    # print(rot_cen_2noised_cen1)
    # print(rot_cen_2noised_cen2)
    rot_cen_2noised_cen = rot_cen_2noised_cen2 @ rot_cen_2noised_cen1
    # 从加噪后的旋转中心到转动后的体素
    noised_cen2noised_ct = np.diag([1, 1, 1, 1]).astype(np.float32)
    noised_cen2noised_ct[:3, 3] = rot_cen
    noised_cen2noised_ct = torch.from_numpy(noised_cen2noised_ct)
    # 从体素到转动后的体素
    ct2noised_ct = noised_cen2noised_ct @ rot_cen_2noised_cen @ ct2rot_cen
    if show_info:
        print("旋转中心\n", rot_cen)
        print("从体素到旋转中心\n", ct2rot_cen)
        print("从旋转中心到加噪后的旋转中心\n", rot_cen_2noised_cen)
        print("从加噪后的旋转中心到转动后的体素\n", noised_cen2noised_ct)
        print("从体素到转动后的体素\n", ct2noised_ct)
    return ct2noised_ct


def _euler_angles2rot_matrix_tensor(theta_x, theta_y, theta_z):
    """
    切记：要求输入为弧度
    :param theta_x: rotation angle along X axis
    :param theta_y: rotation angle along Y axis
    :param theta_z: rotation angle along Z axis
    :return: rotation matrix
    """
    cx = torch.cos(theta_x)
    sx = torch.sin(theta_x)
    rot_x = torch.tensor([[1, 0, 0], [0, cx, sx], [0, -sx, cx]])

    cy = torch.cos(theta_y)
    sy = torch.sin(theta_y)
    rot_y = torch.tensor([[cy, 0, -sy], [0, 1, 0], [sy, 0, cy]])

    cz = torch.cos(theta_z)
    sz = torch.sin(theta_z)
    rot_z = torch.tensor([[cz, sz, 0], [-sz, cz, 0], [0, 0, 1]])

    return torch.matmul(rot_z, torch.matmul(rot_y, rot_x))


def _euler_angles2rot_matrix(theta_x, theta_y, theta_z):
    """
    :param theta_x: rotation angle along X axis
    :param theta_y: rotation angle along Y axis
    :param theta_z: rotation angle along Z axis
    :return: rotation matrix
    """
    cx = np.cos(theta_x)
    sx = np.sin(theta_x)
    rot_x = np.array([[1, 0, 0], [0, cx, sx], [0, -sx, cx]])

    cy = np.cos(theta_y)
    sy = np.sin(theta_y)
    rot_y = np.array([[cy, 0, -sy], [0, 1, 0], [sy, 0, cy]])

    cz = np.cos(theta_z)
    sz = np.sin(theta_z)
    rot_z = np.array([[cz, sz, 0], [-sz, cz, 0], [0, 0, 1]])

    return np.matmul(rot_z, np.matmul(rot_y, rot_x))


def get_default_rt(rot_cen=np.array([89.82421875, 89.82421875, 53]), d_s2c=400,
                   theta=None, beta=None, alpha=None):
    """
    compute the Rt matrix from x-ray source coordinate to CT coordinate

    :param rot_cen: rotation center of C-Arm (mm) in CT coordinate
    :param d_s2c: distance between x-ray source and rot_cen (mm)
    :param theta: rotation angle along z axis of the CT data
    :param beta: angle between z axis of CT and z axis of x-ray source, i.e. rotation angle along y-axis of x-ray source
    :param alpha: rotation angle of image plane along z axis of the C-Arm
    :return: transform matrix from x-ray source coordinate to CT coordinate
    """
    rot_z = _euler_angles2rot_matrix(0, 0, alpha * np.pi / 180)
    # compute rotation matrix from Euler angles (theta_x, theta_y, theta_z) with the default rotation order
    rot_matrix = _euler_angles2rot_matrix(0, np.pi * beta / 180, theta * np.pi / 180) @ rot_z
    # translation vector i.e. the 3D coordinates of the x-ray source in the CT coordinates system
    t_vec = np.array(rot_matrix @ np.mat([0, 0, -d_s2c]).T).squeeze() + np.array(rot_cen)
    return np.column_stack((rot_matrix, t_vec)).astype(np.float32)


def get_rt_tensor(alpha_noise, beta_noise, theta_noise, tx_noise, ty_noise, tz_noise,
                  mode='标准正位', rot_cen=np.array([89.82421875, 89.82421875, 53]), d_s2c=400):
    """
    为了实现vm为损失值的训练，所有的转换过程都得在tensor上进行
    """
    Rt = None
    # 先得到默认的rt
    if mode == "标准正位":
        # 标准正位：（0±5， 270±10， 90±5）, t(0±25， 0±50， 0±25)
        Rt = get_default_rt(alpha=0, beta=270, theta=90,
                            rot_cen=rot_cen, d_s2c=d_s2c)
    Rt = np.vstack((Rt, [0, 0, 0, 1]))
    Rt = torch.from_numpy(Rt)
    noise_rt = add_noise_tensor(torch.tensor([alpha_noise, beta_noise, theta_noise], requires_grad=True),
                                torch.tensor([tx_noise, ty_noise, tz_noise], requires_grad=True),
                                rot_cen=rot_cen)
    Rt = noise_rt @ Rt.to(torch.float)
    return Rt


class VmLoss(nn.Module):
    def __init__(self, Batch_size, voxel_size, interval_num, rot_cen, d_s2c):
        super(VmLoss, self).__init__()
        self.Batch_size = Batch_size
        self.voxel_size = voxel_size
        self.interval_num = interval_num
        # 投影相关参数
        self.rot_cen = rot_cen
        self.d_s2c = d_s2c

    def forward(self, output, target):
        all_vm = torch.tensor(0.0)
        truth_rx, truth_tx = target[:, 0:3], target[:, 3:6]
        pred_rx, pred_tx = output[:, 0: 3], output[:, 3: 6]
        # print(truth_rx.requires_grad)
        # print(pred_rx.requires_grad)
        # 获取标准值
        tru_alpha, tru_beta, tru_theta, tru_tx, tru_ty, tru_tz = label2real(truth_rx, truth_tx, self.Batch_size)
        # 获取预测值
        pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz = label2real(pred_rx, pred_tx, self.Batch_size)
        for i in range(self.Batch_size):
            rt1 = get_rt_tensor(tru_alpha[i], tru_beta[i], tru_theta[i], tru_tx[i], tru_ty[i], tru_tz[i],
                                rot_cen=self.rot_cen, d_s2c=self.d_s2c)
            rt2 = get_rt_tensor(pre_alpha[i], pre_beta[i], pre_theta[i], pre_tx[i], pre_ty[i], pre_tz[i],
                                rot_cen=self.rot_cen, d_s2c=self.d_s2c)
            vm = cal_voxel_mse_loss_tensor(self.voxel_size, self.interval_num, rt1, rt2)
            all_vm += vm
        return all_vm / self.Batch_size
